# -*- coding: utf-8 -*-
import json
import os.path
import shutil
import time
import logging
import paramiko
import textfsm
import argparse
from kafka.producer import KafkaProducer

curr_time = (int(round(time.time() * 1000000)))
curr_path = os.path.dirname(os.path.abspath(__file__))
log_path = os.path.join(curr_path, 'huawei_firewall_session.log')
back_log_path = log_path + str(curr_time)

logging.basicConfig(filename=log_path, filemode='w', level=logging.INFO)


def get_firewall_sessions(hostname, port, username, password):
    client = paramiko.SSHClient()
    client.load_system_host_keys()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    try:
        client.connect(hostname, port, username, password, timeout=2000, allow_agent=False, look_for_keys=False)
        ssh_shell = client.invoke_shell()
        time.sleep(2)
        login_info = ssh_shell.recv(2048).decode('utf-8', 'replace')
        logging.info(login_info)

        ssh_shell.sendall("screen-length 0 temporary\n")
        open_line = ssh_shell.recv(2048).decode('utf-8', 'replace')
        time.sleep(1)
        logging.info(open_line)

        ssh_shell.sendall("display firewall session table verbose\n")
        time.sleep(2)
        output = ''
        while True:
            if ssh_shell.recv_ready():
                output += ssh_shell.recv(4096).decode('utf-8', 'replace')
            else:
                time.sleep(1)
                if not ssh_shell.recv_ready():
                    break
        logging.info(output)
        return output
    except Exception as e:
        print("Failed to connect to the device: {}".format(e))
        return False
    finally:
        client.close()


def send_kafka(device_ip):
    kafka_add = ['10.1.45.6:9092', '10.1.45.7:9092', '10.1.45.8:9092']
    topic_name = 'huawei_firewall_session'
    duanStr = ""
    producer = KafkaProducer(bootstrap_servers=kafka_add, value_serializer=lambda v: json.dumps(v).encode('utf-8'))

    try:
        with open(log_path, 'r') as f:
            while True:
                line = f.readline()
                if not line:
                    break
                if len(line) > 5:
                    # print(len(line))
                    duanStr += line
                else:
                    # print(len(line),'+++',222)
                    # space
                    tf = textfsm.TextFSM(open('/path/to/HuaweiFirewallRegex.template'))
                    # print(duanStr)
                    data = tf.ParseTextToDicts(duanStr)
                    if len(data) > 0:
                        # print(data[0])
                        data[0]['device'] = str(device_ip)
                        data[0]['insert_time'] = str(curr_time)
                        # send kafka
                        producer.send(topic=topic_name, value=data[0])
                        producer.flush()
                        #logging.warning(data[0])
                    duanStr = ""
    finally:
        producer.close()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="The script is get huawei firewalld session rules")
    parser.add_argument("ip", help="Huawei Firewalld Hostname or IP.")
    # parser.add_argument("port", help="Huawei Firewalld Port.")
    # parser.add_argument("username", help="username. Eg: admin???root\!\!\!")
    # parser.add_argument("password", help="password.")
    args = parser.parse_args()
    # session_data = get_firewall_sessions(args.ip, args.port, args.username, args.password)
    session_data = get_firewall_sessions('172.16.48.123', 8888, 'username' + args.ip, 'password')
    if session_data and os.path.isfile(log_path):
        send_kafka(args.ip)
        time.sleep(1)
        try:
            print(log_path, back_log_path)
            shutil.copy(log_path, back_log_path)
            logging.info("copy file success!")
        except IOError as e:
            logging.warning("no copy file")
    else:
        print("Failed to retrieve session data.")
